#!/usr/bin/env python
import os, sys
from argparse import ArgumentParser
import numpy as np
from argparse import ArgumentParser
from astropy.table import Table
import matplotlib.pyplot as plt
from scipy import stats
from scipy.stats import sigmaclip
from astropy.table import Table

def gauss_filter(x,mu,sigma):
    if abs(x)>mu:
        return 1
    return np.exp(-(x-mu)**2/2/sigma**2)+np.exp(-(x+mu)**2/2/sigma**2)
def highpass(x,threshold):
    if abs(x)>threshold:
        return 1
    else:
        return 0

def run(original_scan_res, outdir=None):
    odir = './'
    if outdir is not None:
        odir = outdir + '/'

    tab = Table.read(original_scan_res,format="ipac")[13:63]
    res_x = np.polyfit(np.arange(len(tab["x"].data)), tab["x"].data, deg=1)
    x_fit = res_x[0]*np.arange(len(tab["x"].data))+res_x[1]
    res_y = np.polyfit(np.arange(len(tab["y"].data)), tab["y"].data, deg=1)
    y_fit = res_y[0]*np.arange(len(tab["y"].data))+res_y[1]
    x = tab["x"].data-x_fit
    y = tab["y"].data-y_fit
    flux = tab["flux"].data
    ellip = tab["ellp"].data
    plt.subplot(221)
    plt.scatter(np.arange(len(tab["x"].data)),tab["x"].data,s=10,label='x')
    plt.plot(x_fit)
    plt.title("initial result")
    plt.legend()
    plt.subplot(222)
    plt.scatter(np.arange(len(tab["y"].data)),tab["y"].data,s=10,label='y')
    plt.plot(y_fit)
    plt.legend()
    plt.subplot(223)
    plt.scatter(np.arange(len(tab["flux"].data)),tab["flux"].data,s=10,label='flux')
    plt.legend()
    plt.subplot(224)
    plt.scatter(np.arange(len(tab["ellp"].data)),tab["ellp"].data,s=10,label='ellp')
    plt.legend()
    plt.show() 
#    plt.pause(30)
#    plt.close()

    freqs = np.fft.fftfreq(len(x))
    freqs_shift = np.fft.fftshift(freqs)
    x_trans = np.fft.fft(x)
    x_trans_shift = np.fft.fftshift(x_trans)
    y_trans = np.fft.fft(y)
    y_trans_shift = np.fft.fftshift(y_trans)
    flux_trans = np.fft.fft(flux)
    flux_trans_shift = np.fft.fftshift(flux_trans)
    ellip_trans = np.fft.fft(ellip)
    ellip_trans_shift = np.fft.fftshift(ellip_trans)
    x_trans_shift_mag = np.sqrt((x_trans_shift.real)**2+(x_trans_shift.imag)**2)
    y_trans_shift_mag = np.sqrt((y_trans_shift.real)**2+(y_trans_shift.imag)**2)
    flux_trans_shift_mag = np.sqrt((flux_trans_shift.real)**2+(flux_trans_shift.imag)**2)
    ellip_trans_shift_mag = np.sqrt((ellip_trans_shift.real)**2+(ellip_trans_shift.imag)**2)

    plt.subplot(221)
    plt.plot(freqs_shift,x_trans_shift_mag,label='x')
    plt.axvline(x=1/50)
    plt.axvline(x=1/5)
    plt.yscale("log")
    plt.legend()
    plt.subplot(222)
    plt.plot(freqs_shift,y_trans_shift_mag,label='y')
    plt.axvline(x=1/50)
    plt.axvline(x=1/5)
    plt.yscale("log")
    plt.legend()
    plt.subplot(223)
    plt.plot(freqs_shift,flux_trans_shift_mag,label='flux')
    plt.axvline(x=1/50)
    plt.axvline(x=1/5)
    plt.yscale("log")
    plt.legend()
    plt.subplot(224)
    plt.plot(freqs_shift,ellip_trans_shift_mag,label='ellip')
    plt.axvline(x=1/50)
    plt.axvline(x=1/5)
    plt.yscale("log")
    plt.legend()
    plt.savefig(odir + "fft_result.jpg")
    plt.show()
#    plt.pause(30)
#    plt.close()


    gauss_f = np.zeros(len(freqs_shift))
    highpass_f = gauss_f.copy()
    for i in range(len(freqs_shift)):
        gauss_f[i] = gauss_filter(freqs_shift[i],0.2,0.033)
        highpass_f[i] = highpass(freqs_shift[i],0.2)
    #x_trans_shift_filtered = x_trans_shift*highpass_f
    #y_trans_shift_filtered = y_trans_shift*highpass_f
    x_trans_shift_filtered = x_trans_shift*gauss_f
    y_trans_shift_filtered = y_trans_shift*gauss_f
    flux_trans_shift_filtered = flux_trans_shift*gauss_f
    ellip_trans_shift_filtered = ellip_trans_shift*gauss_f
    x_filtered = np.fft.ifftshift(x_trans_shift_filtered)
    y_filtered = np.fft.ifftshift(y_trans_shift_filtered)
    flux_filtered = np.fft.ifftshift(flux_trans_shift_filtered)
    ellip_filtered = np.fft.ifftshift(ellip_trans_shift_filtered)
    x_filtered = np.fft.ifft(x_filtered).real
    y_filtered = np.fft.ifft(y_filtered).real
    flux_filtered = np.fft.ifft(flux_filtered).real
    ellip_filtered = np.fft.ifft(ellip_filtered).real
    x_trans_shift_filtered_mag = np.sqrt((x_trans_shift_filtered.real)**2+(x_trans_shift_filtered.imag)**2)
    y_trans_shift_filtered_mag = np.sqrt((y_trans_shift_filtered.real)**2+(y_trans_shift_filtered.imag)**2)
    flux_trans_shift_filtered_mag = np.sqrt((flux_trans_shift_filtered.real)**2+(flux_trans_shift_filtered.imag)**2)
    ellip_trans_shift_filtered_mag= np.sqrt((ellip_trans_shift_filtered.real)**2+(ellip_trans_shift_filtered.imag)**2)

    plt.subplot(221)
    plt.plot(freqs_shift,x_trans_shift_filtered_mag,label='x')
    plt.yscale("log")
    plt.title("fft filtered result")
    plt.legend()
    plt.subplot(222)
    plt.plot(freqs_shift,y_trans_shift_filtered_mag,label='y')
    plt.yscale("log")
    plt.legend()
    plt.subplot(223)
    plt.plot(freqs_shift,flux_trans_shift_filtered_mag,label='flux')
    plt.yscale("log")
    plt.legend()
    plt.subplot(224)
    plt.plot(freqs_shift,ellip_trans_shift_filtered_mag,label='ellip')
    plt.yscale("log")
    plt.legend()
    plt.savefig(odir + "fft_filtered_result.jpg")
    plt.show()
#    plt.pause(30)
#    plt.close()


    fig,axes = plt.subplots(2,2,figsize=(16,16))
    axes[0,0].plot(x-np.mean(x),label='x_del')
    axes[0,0].scatter(np.arange(len(x_filtered)),x_filtered-np.mean(x_filtered),s=15)
    axes[0,0].legend()
    axes[0,0].set_ylabel("x_offset(pixel)")
    axes[0,1].plot(y-np.mean(y),label='y_del')
    axes[0,1].scatter(np.arange(len(x_filtered)),y_filtered-np.mean(y_filtered),s=15)
    axes[0,1].legend()
    axes[0,1].set_ylabel("y_offset(pixel)")
    axes[1,0].plot(flux/np.mean(flux),label='flux')
    axes[1,0].scatter(np.arange(len(flux_filtered)),1+flux_filtered/np.mean(flux),s=15)
    axes[1,0].legend()
    axes[1,0].set_ylabel("flux")
    axes[1,1].plot(ellip,label='ellip')
    axes[1,1].scatter(np.arange(len(ellip_filtered)),ellip_filtered+np.mean(ellip),s=15)
    axes[1,1].legend()
    axes[1,1].set_ylabel("ellipcity")
    plt.savefig(odir + "filtered_result.jpg")
    plt.show()
#    plt.pause(30)
#    plt.close()


    d_filtered = np.sqrt((x_filtered-np.mean(x_filtered))**2+(y_filtered-np.mean(y_filtered))**2)
    d_sigma_clip,_,_ = sigmaclip(d_filtered,5,5)
    flux_sigma_clip,_,_ = sigmaclip(flux_filtered,5,5)
    ellip_sigma_clip,_,_ = sigmaclip(ellip_filtered,5,5)
    print(ellip_sigma_clip)
    print("d_sig",np.std(d_sigma_clip))
    print("flux_sig",np.std(flux_sigma_clip)/np.mean(flux))
    print("ellip_mean",np.mean(ellip),"ellip_sig",np.std(ellip_sigma_clip))
    tab = Table()
    tab['fluxerr'] = np.array([np.std(flux_sigma_clip)/np.mean(flux)])*100
    tab['fluxerr'].unit = '%'
    tab['poserr'] = np.array([np.std(d_sigma_clip)])*10
    tab['poserr'].unit = 'um'
    tab['ellp'] = np.array([np.mean(ellip)])
    tab['ellperr'] = np.array([np.std(ellip_sigma_clip)])
    tab.write(odir + "scan_filtered_result.tab",format='ipac',overwrite=True)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--scan_res', '-s', type=str,
        help='original scan spot data results/data/scan_*_data.dat')
    parser.add_argument('--outdir', '-o', default='./', type=str,
        help='outdir of filtered spot scan results/figures')

    args = parser.parse_args()
    run(original_scan_res=args.scan_res, outdir=args.outdir)
